from openai import OpenAI
import os

# 从环境变量中获取 DeepSeek API 密钥
DeepSeek_API_KEY = os.getenv("DeepSeek_API_KEY")

# 创建 OpenAI 客户端，使用 DeepSeek API 密钥和基础 URL
client = OpenAI(api_key=DeepSeek_API_KEY, base_url="https://api.deepseek.com")


# 定义一个函数来获取模型的响应
def get_response(messages):
    # 使用客户端的 chat.completions.create 方法生成响应
    response = client.chat.completions.create(
        model="deepseek-chat",  # 使用的模型名称
        messages=messages,  # 传入的消息列表
        max_tokens=1024,  # 最大生成的 token 数量
        temperature=0.7,  # 生成响应的随机性
        stream=True  # 是否以流的方式返回响应
    )
    return response


# 主函数，用于与用户进行交互
def main():
    # 初始化消息列表，包含一个系统消息
    messages = [
        {"role": "system", "content": "You are a helpful assistant"},
    ]

    # 欢迎信息
    print("欢迎使用聊天助手！输入 'exit' 退出。")

    # 持续接收用户输入的循环
    while True:
        # 获取用户输入
        user_input = input("你: ")

        # 如果用户输入 'exit'，退出循环
        if user_input.lower() == 'exit':
            print("再见！")
            break

        # 将用户输入添加到消息列表中
        messages.append({"role": "user", "content": user_input})

        # 获取模型的响应
        response = get_response(messages)

        # 初始化助手的响应字符串
        assistant_response = ""

        # 逐块处理模型的响应
        for chunk in response:
            # 如果 chunk 中有内容，添加到助手的响应字符串中并打印
            if chunk.choices[0].delta.content:
                assistant_response += chunk.choices[0].delta.content
                print(chunk.choices[0].delta.content, end="")

        # 将助手的响应添加到消息列表中
        messages.append({"role": "assistant", "content": assistant_response})

        # 换行以便更好地阅读
        print()


# 确保 main 函数在脚本直接运行时执行
if __name__ == "__main__":
    main()
